﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.InteropServices;
using RayDen.Library.Components.SystemComponents;
using RayDen.Library.Core;
using RayDen.Library.Core.Primitives;
using RayDen.Library.Entity.Scene;
using RayDen.RayEngine.Core.Interface;
using RayDen.RayEngine.Core.Types;
using RayDen.RayEngine.Data;

using Triangle=RayDen.Library.Core.Primitives.TriangleDataInfo;

namespace RayDen.RayEngine.IntersectionEngines
{
    public class TopDownBvhAccell
    {
        private const int PrimitivesInNode = 4;
        public int TraversalStackSize = 64;
        public const uint NoHit = RayBuffer.NotHit;

        private int costSamples = 4, isectCost = 80, traversalCost = 10;
        private float emptyBonus = 0.25f;

        public Node RootNode;

     
        public class Node
        {
            public AABB Bound;
            public Node Left, Right;
            public List<Triangle> Primitives;
            public int PrimitiveIndexStart;
        }

        private Point[] sceneVertices;
        public Triangle[] triangles;
        private List<ITriangleMesh> meshes;
        public List<Triangle> sortedPrims = new List<Triangle>();

        public GpuLBVHNode[] nodes;


        private void LinearizeBVH()
        {
            nodes = new GpuLBVHNode[bvhNodesCount];
            //int offset = 0;

            this.BuildArray(RootNode, 0);
            //this.FlattenBVHTree(RootNode, ref offset);
        }

        private int BuildArray(Node node, int offset)
        {
            var nd = node;
            while (nd != null)
            {
                var firstOffset = offset;
                var p = new GpuLBVHNode()
                {
                    bbox = new BBox(nd.Bound),
                    PrimitiveIndex = nd.Primitives != null ? (uint) nd.PrimitiveIndexStart : NoHit,
                };
                offset = BuildArray(nd.Left, offset + 1);
                if (nd.Primitives == null)
                {
                    p.SkipIndex = (uint) offset;
                    p.Axis = (byte) nd.Bound.MaximumExtent();
                }
                else
                {
                    p.count = (byte) nd.Primitives.Count;
                }
                nodes[firstOffset] = p;
                nd = nd.Right;
            }
            return offset;
        }


        public void BuildTree(List<Triangle> primitives, int start, int end)
        {
            Tracer.TraceLine("Building accelleration {0} triangles", primitives.Count);

            if (start - end == 0)
                return;
            AABB bound = primitives.GetRange(start, end - start).Select(prim => prim.WorldBound(sceneVertices)).Aggregate((b1, b2) => b1.Union(b2));
            RootNode = new Node() { Bound = bound };
            bvhNodesCount++;
            if ((end - start) <= PrimitivesInNode)
            {
                RootNode.Primitives = primitives.GetRange(start, (end - start));
                RootNode.PrimitiveIndexStart = start;
                return;
            }



            RootNode.Left = this.BuildSubTree(primitives, start, (start + end) / 2);
            RootNode.Right = this.BuildSubTree(primitives, 1+(start + end) / 2, end);

            Tracer.TraceLine("Linearize BVH");
            this.LinearizeBVH();
            Tracer.TraceLine("Complete");

        }

        private Node BuildSubTree(List<Triangle> primitives, int start, int end)
        {
            if (start - end == 0)
                return null;
            AABB bound = primitives.GetRange(start, end - start).Select(prim => prim.WorldBound(sceneVertices)).Aggregate((b1, b2) => b1.Union(b2));
            var node = new Node() { Bound = bound, Primitives =  new List<TriangleDataInfo>()};
            bvhNodesCount++;
            if ((end - start) <= PrimitivesInNode)
            {

                var firstOffset = sortedPrims.Count;
                for (var i = start; i < end; ++i)
                {
                    var primNum = i;
                    sortedPrims.Add(triangles[primNum]);
                }
                node.PrimitiveIndexStart = firstOffset;
                node.Primitives = primitives.GetRange(start, end - start);
                return node;
            }
            var centroidBounds = new AABB();
            for (var i = start; i < end; ++i)
                centroidBounds = centroidBounds.Union(primitives[i].WorldBound(sceneVertices).Center);
            int dim = centroidBounds.MaximumExtent();

            if ((MathLab.SmallEnough(centroidBounds.Min[dim], centroidBounds.Max[dim])))
            {
                var firstOffset = start;
                for (int i = start, ii = 0; i < end; ++i, ii++)
                {
                    var primNum = i;
                    sortedPrims.Add(triangles[primNum]);
                }

                node.PrimitiveIndexStart = firstOffset;
                node.Primitives = primitives.GetRange(start, end - start);

                return node;
            }

            var mid = (start + end) / 2;
            node.Left = this.BuildSubTree(primitives, start, mid);
            node.Right = this.BuildSubTree(primitives, mid+1, end);
            return node;
        }



        #region IAccellerationStructure Members

        private int bvhNodesCount;
        public void Init(RayEngineScene scene, TriangleMesh[] geometry)
        {
            var sceneModel = scene;
            meshes = new List<ITriangleMesh>();
            meshes.AddRange(scene.Meshes);
            this.sceneVertices = sceneModel.Vertices;
            var trianglesCache = scene.Triangles.ToList();
            Comparison<Triangle> tc = (t1, t2) => CompareTriangles(t1, t2);
            //trianglesCache.Sort(tc);
            this.triangles = trianglesCache.ToArray();
            this.BuildTree(trianglesCache, 0, trianglesCache.Count);
            Tracer.TraceLine("Tree is builded, {0} nodes total", bvhNodesCount);
        }

        public int CompareTriangles(Triangle t1, Triangle t2, Axis a = Axis.PLANE_Y)
        {
            var p1 = t1.Center(sceneVertices);
            var p2 = t2.Center(sceneVertices);
            return p1.Component(a).CompareTo(p2.Component(a));
        }
        #endregion

        public IEnumerable<GpuLBVHNode> GetData()
        {
            return this.nodes;
        }

        void FindBestSplit(List<Triangle> list, int begin, int end, out float splitValue, out int bestAxis)
        {
            splitValue = 0;
            if (end - begin == 2)
            {
                splitValue = (list[begin].WorldBound(sceneVertices).Max[0] + list[begin].WorldBound(sceneVertices).Min[0] + list[end - 1].WorldBound(sceneVertices).Max[0] + list[end - 1].WorldBound(sceneVertices).Min[0]) / 2f;
                bestAxis = 0;
            }
            else
            {
                Vector mean2 = new Vector(), var = new Vector();

                for (int i = begin; i < end; i++)
                    mean2 += (Vector) (list[i].WorldBound(sceneVertices).Max + list[i].WorldBound(sceneVertices).Min);
                mean2 /= end - begin;

                for (int i = begin; i < end; i++)
                {
                    Vector v = list[i].WorldBound(sceneVertices).Max + list[i].WorldBound(sceneVertices).Min - mean2;
                    v.x *= v.x;
                    v.y *= v.y;
                    v.z *= v.z;
                    var += v;
                }
                // Select axis with more variance
                if (var.x > var.y && var.x > var.z)
                    bestAxis = 0;
                else if (var.y > var.z)
                    bestAxis = 1;
                else
                    bestAxis = 2;

                if (costSamples > 1)
                {
                    AABB nodeBounds = new AABB();//!!!!
                    for (int i = begin; i < end; i++)
                        nodeBounds = nodeBounds.Union(list[i].WorldBound(sceneVertices));

                    Vector d = nodeBounds.Max - nodeBounds.Min;
                    float totalSA = (2f * (d.x * d.y + d.y * d.z + d.z * d.x));
                    float invTotalSA = 1f / totalSA;

                    // Sample cost for split at some points
                    float increment = 2f * d[bestAxis] / (costSamples + 1);
                    float bestCost = float.PositiveInfinity;
                    for (float splitVal = 2f * nodeBounds.Min[bestAxis] + increment; splitVal < 2f * nodeBounds.Max[bestAxis]; splitVal += increment)
                    {
                        int nBelow = 0, nAbove = 0;
                        AABB bbBelow = new AABB(Point.Zero), bbAbove = new AABB(Point.Zero);
                        for (int j = begin; j < end; j++)
                        {
                            if ((list[j].WorldBound(sceneVertices).Max[bestAxis] + list[j].WorldBound(sceneVertices).Min[bestAxis]) < splitVal)
                            {
                                nBelow++;
                                bbBelow = bbBelow.Union(list[j].WorldBound(sceneVertices));
                            }
                            else
                            {
                                nAbove++;
                                bbAbove = bbAbove.Union(list[j].WorldBound(sceneVertices));
                            }
                        }
                        Vector dBelow = bbBelow.Max - bbBelow.Min;
                        Vector dAbove = bbAbove.Max - bbAbove.Min;
                        float belowSA = 2f * ((dBelow.x * dBelow.y + dBelow.x * dBelow.z + dBelow.y * dBelow.z));
                        float aboveSA = 2f * ((dAbove.x * dAbove.y + dAbove.x * dAbove.z + dAbove.y * dAbove.z));
                        float pBelow = belowSA * invTotalSA;
                        float pAbove = aboveSA * invTotalSA;
                        float eb = (nAbove == 0 || nBelow == 0) ? emptyBonus : 0;
                        float cost = traversalCost + isectCost * (1f - eb) * (pBelow * nBelow + pAbove * nAbove);
                        // Update best split if this is lowest cost so far
                        if (cost < bestCost)
                        {
                            bestCost = cost;
                            splitValue = splitVal;
                        }
                    }
                }
                else
                {
                    // Split in half around the mean center
                    splitValue = mean2[bestAxis];
                }
            }
        }


    }
}
